package org.usmile.algorithms.leetcode.middle;

import org.usmile.algorithms.leetcode.TreeNode;

import java.util.HashMap;
import java.util.Map;

/**
 * 437. 路径总和 III
 *
 * 给定一个二叉树的根节点 root ，和一个整数 targetSum ，求该二叉树里节点值之和等于 targetSum 的 路径 的数目。
 * 路径 不需要从根节点开始，也不需要在叶子节点结束，但是路径方向必须是向下的（只能从父节点到子节点）。
 *
 * 示例 1：
 * 输入：root = [10,5,-3,3,2,null,11,3,-2,null,1], targetSum = 8
 * 输出：3
 * 解释：和等于 8 的路径有 3 条，如图所示。
 *
 * 示例 2：
 * 输入：root = [5,4,8,11,null,13,4,7,2,null,null,5,1], targetSum = 22
 * 输出：3
 *
 * 提示:
 * 二叉树的节点个数的范围是 [0,1000]
 * -109 <= Node.val <= 109
 * -1000 <= targetSum <= 1000
 */
public class _0437 {
}

class _0437_Solution {
    public int pathSum(TreeNode root, int sum) {
        Map<Long, Integer> prefixSumMap = new HashMap<>();
        int currSum = 0;
        prefixSumMap.put(0L, 1);

        return dfs(root, currSum, sum, prefixSumMap);
    }

    private int dfs(TreeNode node,
                     long currSum, int target,
                     Map<Long, Integer> prefixSumMap) {
        if (node == null) {
            return 0;
        }

        currSum += node.val;
        // pre_i - pre_j = targetSum pre_i为当前结点前缀和，pre_j为当前路径上i节点之前的结点j的前缀和
        int count = prefixSumMap.getOrDefault(currSum - target, 0);
        prefixSumMap.put(currSum, prefixSumMap.getOrDefault(currSum, 0) + 1);

        count += dfs(node.left, currSum, target, prefixSumMap);
        count += dfs(node.right, currSum, target, prefixSumMap);

        prefixSumMap.put(currSum, prefixSumMap.get(currSum) - 1);

        return count;
    }
}
